import os.path
import numpy as np
import torch
from copy import deepcopy
import wandb
from medpy import metric
from tqdm import tqdm
import torch.nn as nn
import random
from scipy.ndimage import zoom
from utils import loadyaml, _get_logger, mk_path, get_current_consistency_weight, DiceLoss, update_ema_variables
from utils import build_lr_scheduler, build_optimizer
from model import build_model
from datasets import build_loader


# paper link https://arxiv.org/abs/2202.00677

def main():
    path = r"config/ict-medseg_unet_30k_224x224_ACDC_Semi.yaml"
    root = os.path.dirname(os.path.realpath(__file__))  # 获取绝对路径
    args = loadyaml(os.path.join(root, path))  # 加载yaml
    if args.cuda:
        args.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    else:
        args.device = torch.device("cpu")

    root = os.path.dirname(os.path.realpath(__file__))  # 获取绝对路径
    args.save_path = os.path.join(root, args.save_path)
    mk_path(args.save_path)  # 创建文件保存位置
    # 创建 tensorboardX日志保存位置
    mk_path(os.path.join(args.save_path, "tensorboardX"))
    mk_path(os.path.join(args.save_path, "model"))  # 创建模型保存位置
    args.model_save_path = os.path.join(args.save_path, "model", "model.pth")
    args.ema_model_save_path = os.path.join(args.save_path, "model", "ema_model_model.pth")

    wandb.init(
        entity="jokerak777",  # wandb上对应的team名称（必填）,类似于用户名
        project="semi",  # wandb上对应的team名称（必填）,
        name="test",  # 本次实验的名称（可选，如果不设置，wandb会自动生成本次实验名称）
        config=args,
        dir=os.path.join(args.save_path, "tensorboardX")
    )
    args.logger = _get_logger(os.path.join(args.save_path, "log.log"), "info")
    args.tqdm = os.path.join(args.save_path, "tqdm.log")
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

    torch.backends.cudnn.deterministic = False  # 单卡的不需要分布式
    torch.backends.cudnn.benchmark = True  # 寻找最佳 的训练路径

    trainloader, test_loader = build_loader(args)  # 构建数据集
    args.epochs = args.total_itrs // args.step_size  # 设置模型epoch
    args.logger.info("==========> trainloader length:{}".format(len(trainloader.dataset)))
    args.logger.info("==========> trainloader length:{}".format(len(trainloader)))
    args.logger.info("==========> test_dataloader length:{}".format(len(test_loader.dataset)))
    args.logger.info("==========> epochs length:{}".format(args.epochs))

    # step 1: 构建模型
    model = build_model(args=args).to(device=args.device)  # 创建模型
    # ema_model = deepcopy(model)  # 创建ema_model
    ema_model = build_model(args=args).to(device=args.device)
    for name, param in ema_model.named_parameters():
        param.requires_grad = False

    wandb.watch(model, log='all', log_freq=1000)

    # step 2: 训练模型
    ICT_MedSeg(model, ema_model, trainloader, test_loader, args)


def ICT_MedSeg(model, ema_model, trainloader, test_loader, args):
    optimizer = build_optimizer(args=args, model=model)
    lr_scheduler = build_lr_scheduler(args=args, optimizer=optimizer)
    # config network and criterion
    ce_loss = nn.CrossEntropyLoss(ignore_index=255)
    dice_loss = DiceLoss(args.num_classes)

    model.train()
    cur_itrs = 0
    train_loss = 0
    best_dice = 0.0
    best_ema_dice = 0.0

    max_epoch = args.total_itrs // len(trainloader)
    pbar = tqdm(total=args.total_itrs)
    args.logger.info("==========> max_epoch:{}".format(max_epoch))

    for epoch_num in range(max_epoch):
        for sampled_batch in trainloader:
            pbar.update(1)
            cur_itrs += 1
            volume_batch, label_batch = sampled_batch
            volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda()
            unlabeled_volume_batch = volume_batch[args.batch_size:]  # 无标注数据
            labeled_volume_batch = volume_batch[:args.batch_size]  # 有标注数据

            # ICT mix factors
            ict_mix_factors = np.random.beta(args.ict_alpha, args.ict_alpha, size=(args.unlabel_batch_size//2, 1, 1, 1))  # 设置mixup随机混合参数
            ict_mix_factors = torch.tensor(ict_mix_factors, dtype=torch.float).cuda()
            unlabeled_volume_batch_0 = unlabeled_volume_batch[0:args.unlabel_batch_size//2, ...]
            unlabeled_volume_batch_1 = unlabeled_volume_batch[args.unlabel_batch_size//2:, ...]

            # Mix images
            batch_ux_mixed = unlabeled_volume_batch_0 * (1.0 - ict_mix_factors) + unlabeled_volume_batch_1 * ict_mix_factors
            input_volume_batch = torch.cat([labeled_volume_batch, batch_ux_mixed], dim=0)
            outputs = model(input_volume_batch)
            outputs_soft = torch.softmax(outputs, dim=1)
            with torch.no_grad():
                ema_output_ux0 = torch.softmax(ema_model(unlabeled_volume_batch_0), dim=1)
                ema_output_ux1 = torch.softmax(ema_model(unlabeled_volume_batch_1), dim=1)
                batch_pred_mixed = ema_output_ux0 * (1.0 - ict_mix_factors) + ema_output_ux1 * ict_mix_factors

            loss_ce = ce_loss(outputs[:args.batch_size], label_batch[:args.batch_size][:].long())
            loss_dice = dice_loss(outputs_soft[:args.batch_size], label_batch[:args.batch_size].unsqueeze(1))

            supervised_loss = 0.5 * (loss_dice + loss_ce)
            consistency_weight = get_current_consistency_weight(epoch=cur_itrs // 150, args=args)
            consistency_loss = torch.mean((outputs_soft[args.batch_size:] - batch_pred_mixed) ** 2)

            loss = supervised_loss + consistency_weight * consistency_loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            update_ema_variables(model, ema_model, args.ema_decay, cur_itrs)
            lr = args.lr * (1.0 - cur_itrs / args.total_itrs) ** 0.9
            for param_group in optimizer.param_groups:
                param_group['lr'] = lr

            wandb.log({
                "ICT_MedSeg/loss": loss.item(),
                "ICT_MedSeg/lr": lr,
                "ICT_MedSeg/consistency_weight": consistency_weight,
                "ICT_MedSeg/consistency_loss": consistency_loss
            })
            if cur_itrs % args.step_size == 0:
                dice, hd95 = test_acdc(model=model, test_loader=test_loader, args=args, name="test_model")

                wandb.log({
                    'model_dice': dice,
                    'model_hd95': hd95,
                })
                args.logger.info("val [{}/{} ({:.0f}%)]\t model dice:{:.4f} \t hd95:{:.4f} ".format(cur_itrs, args.total_itrs,
                                                                                                    100. * cur_itrs / args.total_itrs,
                                                                                                    dice, hd95))

                if dice > best_dice:
                    best_dice = dice
                    torch.save({
                        "cur_itrs": cur_itrs,
                        "best_dice": best_dice,
                        "model": model.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "lr_scheduler": lr_scheduler.state_dict(),
                    }, args.model_save_path)

                dice, hd95 = test_acdc(model=ema_model, test_loader=test_loader, args=args, name="test_ema_model")

                if dice > best_ema_dice:
                    best_ema_dice = dice
                    torch.save({
                        "cur_itrs": cur_itrs,
                        "best_dice": best_dice,
                        "model": ema_model.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "lr_scheduler": lr_scheduler.state_dict(),
                    }, args.ema_model_save_path)

                wandb.log({
                    'ema_model_dice': dice,
                    'ema_model_hd95': hd95,
                })

                args.logger.info("val [{}/{} ({:.0f}%)]\t ema_model dice:{:.4f} \t hd95:{:.4f} ".format(cur_itrs, args.total_itrs,
                                                                                                        100. * cur_itrs / args.total_itrs,
                                                                                                        dice, hd95))

                args.logger.info("best model dice:{:.4f} ,best ema_model dice {:.4f}".format(best_dice, best_ema_dice))

                model.train()
                ema_model.train()

            if cur_itrs > args.total_itrs:
                return

    # do this after training
    wandb.finish()


def test_acdc(model, test_loader, args, name):

    model.eval()
    metric_list = 0.0
    for i_batch, sampled_batch in enumerate(test_loader):
        image = sampled_batch[0].to(args.device)
        label = sampled_batch[1].to(args.device)
        metric_i = test_single_volume(image, label, model, classes=args.num_classes, patch_size=args.test_crop_size)
        metric_list += np.array(metric_i)

        if i_batch == 0:
            slice = image[0, 0, :, :].cpu().detach().numpy()
            x, y = slice.shape[0], slice.shape[1]
            slice = zoom(slice, (args.test_crop_size[0] / x, args.test_crop_size[1] / y), order=0)
            img = torch.from_numpy(slice).unsqueeze(0).unsqueeze(0).float().cuda()
            label_pred = torch.argmax(torch.softmax(model(img), dim=1), dim=1, keepdim=False).squeeze(0)
            label_pred = label_pred.cpu().detach().numpy()
            label_pred = zoom(label_pred, (x / args.test_crop_size[0], y / args.test_crop_size[1]), order=0)
            label_pred = test_loader.dataset.label_to_img(label_pred)

            label_true = label[0, 0, :, :].squeeze().cpu().detach().numpy()
            label_true = test_loader.dataset.label_to_img(label_true)

            Img = wandb.Image(img[0])
            label_pred = wandb.Image(label_pred)
            label_true = wandb.Image(label_true)

            wandb.log({"{}/Image".format(name): Img, "{}/label_pred".format(name): label_pred, "{}/label_true".format(name): label_true})

    metric_list = metric_list / len(test_loader.dataset)
    performance2 = np.mean(metric_list, axis=0)[0]
    mean_hd952 = np.mean(metric_list, axis=0)[1]
    return performance2, mean_hd952


def test_single_volume(image, label, net, classes, patch_size=[256, 256]):
    image, label = image.squeeze(0).cpu().detach().numpy(), label.squeeze(0).cpu().detach().numpy()
    prediction = np.zeros_like(label)
    for ind in range(image.shape[0]):
        slice = image[ind, :, :]
        x, y = slice.shape[0], slice.shape[1]
        slice = zoom(slice, (patch_size[0] / x, patch_size[1] / y), order=0)
        input = torch.from_numpy(slice).unsqueeze(0).unsqueeze(0).float().cuda()
        net.eval()
        with torch.no_grad():
            out = torch.argmax(torch.softmax(net(input), dim=1), dim=1).squeeze(0)
            out = out.cpu().detach().numpy()
            pred = zoom(out, (x / patch_size[0], y / patch_size[1]), order=0)
            prediction[ind] = pred
    metric_list = []
    for i in range(1, classes):
        metric_list.append(calculate_metric_percase(prediction == i, label == i))
    return metric_list


def calculate_metric_percase(pred, gt):
    pred[pred > 0] = 1
    gt[gt > 0] = 1
    if pred.sum() > 0:
        dice = metric.binary.dc(pred, gt)
        hd95 = metric.binary.hd95(pred, gt)
        return dice, hd95

    elif pred.sum() > 0 and gt.sum() == 0:
        return 1, 0
    else:
        return 0, 0


if __name__ == '__main__':
    main()
